from flask import Blueprint, render_template, request, session, redirect, url_for, jsonify, current_app, g
import requests
from datetime import datetime
import json
import re
from bs4 import BeautifulSoup
import os
import time
from urllib.parse import quote_plus
from flask_login import login_required, current_user

# 创建蓝图
ai_assistant_bp = Blueprint('ai_assistant', __name__, url_prefix='/ai_assistant')

# 聊天历史记录存储
# 使用用户名作为键，值为消息列表
chat_histories = {}

# AI助手页面
@ai_assistant_bp.route('/')
@login_required
def index():
    # 初始化用户的聊天历史
    username = current_user.username
    if username not in chat_histories:
        chat_histories[username] = []
    
    current_time = datetime.now().strftime('%H:%M:%S')
    return render_template('ai_assistant.html', username=username, current_time=current_time)

# 清除历史记录API
@ai_assistant_bp.route('/clear_history', methods=['POST'])
def clear_history():
    session_id = session.get('user_id', 'anonymous')
    
    if 'chat_history' in g and session_id in g.chat_history:
        g.chat_history[session_id] = []
    
    return jsonify({'status': 'success', 'message': '历史记录已清空'})

# 提取可能需要实时信息的搜索查询
def extract_search_queries(message):
    """
    根据用户消息提取可能需要实时网络信息的搜索查询
    返回一个查询列表或None
    """
    # 下面是一些可能需要实时信息的关键词模式
    time_sensitive_patterns = [
        r'最近|最新|近期|今日|昨日|本周|本月|今年',
        r'现在|当前|目前|实时',
        r'新闻|事件|热点|趋势|动态',
        r'多少钱|价格|股价|行情',
        r'天气|温度|气温|降雨|台风',
        r'(什么|谁|哪个)是.*?最(好|新|热门|受欢迎)',
        r'发生了什么',
        r'如何.*?(现在|当前|今日)',
        r'疫情|肺炎|新冠',
        r'选举|投票|政策|法规'
    ]
    
    # 尝试匹配时间敏感的模式
    for pattern in time_sensitive_patterns:
        if re.search(pattern, message):
            # 如果消息较长，尝试提取关键信息
            if len(message) > 50:
                # 移除常见的询问词和停用词
                cleaned_message = re.sub(r'^(请问|问一下|我想知道|能告诉我|帮我查一下|查询|搜索)', '', message)
                # 保留最后一个问号前的内容
                if '？' in cleaned_message:
                    cleaned_message = cleaned_message.split('？')[0] + '？'
                elif '?' in cleaned_message:
                    cleaned_message = cleaned_message.split('?')[0] + '?'
                
                return [cleaned_message.strip()]
            else:
                return [message]
    
    # 匹配直接询问特定事实的模式
    fact_patterns = [
        r'(谁|什么|哪里|哪个|何时|如何|为什么)是',
        r'(谁|什么|哪里|哪个|何时|如何|为什么)(创建|发明|开发)',
        r'.*?(的人口|的首都|的总统|的CEO|的创始人)',
        r'.*?的(地址|位置|电话|评价)',
        r'如何(使用|安装|配置|设置)',
        r'(发生|开始|结束)于(何时|哪一年)'
    ]
    
    for pattern in fact_patterns:
        if re.search(pattern, message):
            # 对于特定事实查询，直接使用原消息
            return [message]
    
    # 没有匹配到任何需要实时信息的模式
    return None

# 网络搜索功能
def search_web(query):
    """
    使用Google搜索API查询信息
    返回搜索结果列表或None
    """
    try:
        # 获取API配置
        api_key = current_app.config.get('GOOGLE_API_KEY', '')
        search_engine_id = current_app.config.get('SEARCH_ENGINE_ID', '')
        api_url = current_app.config.get('GOOGLE_CSE_URL', 'https://www.googleapis.com/customsearch/v1')
        
        if not api_key or not search_engine_id:
            print("Google搜索API未配置")
            return None
        
        # 构建搜索URL
        encoded_query = quote_plus(query)
        url = f"{api_url}?key={api_key}&cx={search_engine_id}&q={encoded_query}"
        
        # 发送请求
        response = requests.get(url, timeout=10)
        
        if response.status_code == 200:
            data = response.json()
            
            if 'items' in data and len(data['items']) > 0:
                results = []
                
                # 提取前5个搜索结果
                for item in data['items'][:5]:
                    results.append({
                        'title': item.get('title', ''),
                        'link': item.get('link', ''),
                        'snippet': item.get('snippet', '')
                    })
                
                return results
        
        # 备用搜索方法：如果Google API不可用，尝试使用其他方法
        return fallback_search(query)
    
    except Exception as e:
        print(f"搜索时出错: {str(e)}")
        # 出错时尝试备用搜索
        return fallback_search(query)

# 备用搜索方法
def fallback_search(query):
    """
    当主要搜索方法失败时的备用搜索实现
    通过直接抓取搜索引擎结果
    """
    try:
        # 使用DuckDuckGo lite版本，格式较简单，容易解析
        headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
        }
        
        encoded_query = quote_plus(query)
        url = f"https://lite.duckduckgo.com/lite/?q={encoded_query}"
        
        response = requests.get(url, headers=headers, timeout=10)
        
        if response.status_code == 200:
            soup = BeautifulSoup(response.text, 'html.parser')
            results = []
            
            # DuckDuckGo lite版本中，结果在表格中
            for tr in soup.find_all('tr'):
                if len(tr.find_all('td')) >= 2:
                    tds = tr.find_all('td')
                    
                    # 第一个td包含链接和标题
                    link_element = tds[0].find('a')
                    if link_element and link_element.get('href'):
                        link = link_element.get('href')
                        title = link_element.text.strip()
                        
                        # 第二个td通常包含简介
                        snippet = tds[1].text.strip() if len(tds) > 1 else ""
                        
                        results.append({
                            'title': title,
                            'link': link,
                            'snippet': snippet
                        })
                        
                        # 只取前5个结果
                        if len(results) >= 5:
                            break
            
            if results:
                return results
        
        # 如果抓取失败，返回空结果
        return None
    
    except Exception as e:
        print(f"备用搜索失败: {str(e)}")
        return None

# AI聊天API
@ai_assistant_bp.route('/chat', methods=['POST'])
def ai_chat():
    data = request.json
    user_message = data.get('message', '')
    with_memory = data.get('with_memory', True)
    use_web_search = data.get('use_web_search', False)

    print(f"user_message: {user_message}")
    print(f"with_memory: {with_memory}")
    print(f"use_web_search: {use_web_search}")
    
    image_data = None
    if 'image' in data:
        image_data = data['image']
        print(f"image_data: {image_data}")
    
    session_id = session.get('user_id', 'anonymous')
    
    print(f"session_id: {session_id}")

    # 确保chat_history字典存在
    if 'chat_history' not in g:
        g.chat_history = {}
    
    # 确保该用户的聊天历史存在
    if session_id not in g.chat_history:
        g.chat_history[session_id] = []
    
    # 将用户消息添加到历史记录
    if image_data:
        g.chat_history[session_id].append({'role': 'user', 'content': user_message, 'has_image': True})
        
    else:
        g.chat_history[session_id].append({'role': 'user', 'content': user_message})
    
    # 准备请求数据
    messages = []
    
    # 如果启用了记忆功能，使用历史聊天记录
    if with_memory:
        # 复制用户的聊天历史用于请求
        for msg in g.chat_history[session_id]:
            if msg.get('has_image', False):
                if 'image_url' in msg:
                    # 如果历史消息中有图片URL，发送它
                    messages.append({
                        'role': msg['role'],
                        'content': [
                            {'type': 'text', 'text': msg['content']},
                            {'type': 'image_url', 'image_url': {'url': msg['image_url']}}
                        ]
                    })
                else:
                    # 如果是历史图片消息但没有URL，只发送文本
                    messages.append({'role': msg['role'], 'content': msg['content']})
            else:
                # 纯文本消息
                messages.append({'role': msg['role'], 'content': msg['content']})
    else:
        # 不使用记忆功能，只发送系统提示和当前消息
        messages.append({'role': 'system', 'content': '你是一个有用的AI助手。'})
        if image_data:
            # 当前请求有图片
            messages.append({
                'role': 'user',
                'content': [
                    {'type': 'text', 'text': user_message},
                    {'type': 'image_url', 'image_url': {'url': image_data}}
                ]
            })
        else:
            # 当前请求没有图片
            messages.append({'role': 'user', 'content': user_message})
    
    # 如果请求中包含图片，处理图片数据
    if image_data and 'has_image' in g.chat_history[session_id][-1]:
        # 将Base64图片数据直接存储在最后一条消息中
        g.chat_history[session_id][-1]['image_url'] = image_data
    
    # 检查是否需要进行网络搜索
    # 当用户启用联网功能并且消息可能需要实时信息时进行搜索
    web_search_results = None
    if use_web_search and user_message:
        search_queries = extract_search_queries(user_message)
        if search_queries:
            web_search_results = []
            for query in search_queries:
                search_result = search_web(query)
                if search_result:
                    web_search_results.append({
                        'query': query,
                        'results': search_result
                    })
    
    try:
        # 获取API配置
        api_key = current_app.config.get('GEMINI_API_KEY', '')
        api_base = current_app.config.get('GEMINI_API_URL', '')
        api_model = current_app.config.get('GEMINI_MODEL', 'gemini-2.0-flash-lite')
        
        # 准备请求头
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {api_key}"
        }
        
        # 如果有网络搜索结果，将其添加到消息中
        if web_search_results:
            search_info = "以下是根据您的问题进行的网络搜索结果：\n\n"
            for result in web_search_results:
                search_info += f"查询: {result['query']}\n"
                for item in result['results']:
                    search_info += f"- {item['title']}: {item['snippet']}\n  来源: {item['link']}\n\n"
            
            search_info += "请基于这些搜索结果和您的知识回答用户的问题。如果使用了搜索结果中的信息，请在回答末尾添加：*(该回答使用了联网搜索功能提供的信息)*"
            
            # 添加系统消息包含搜索结果
            messages.append({
                'role': 'system',
                'content': search_info
            })

        print(f"messages: {messages}")
        # 准备请求数据
        request_data = {
            "model": api_model,
            "messages": messages,
            "temperature": 0.7,
            "max_tokens": 1500
        }
        
        # 如果是包含图片的请求，使用vision模型
        # if image_data and not with_memory:
        #     # request_data["model"] = "gpt-4-vision-preview"
        #     request_data["max_tokens"] = 1000
        
        print(f"请求数据: {request_data}")
        # 发送请求到API
        response = requests.post(
            f"{api_base}/chat/completions",
            headers=headers,
            json=request_data,
            timeout=60
        )
        
        response_json = response.json()
        
        # 如果响应中有 choices，提取内容
        if 'choices' in response_json and len(response_json['choices']) > 0:
            ai_response = response_json['choices'][0]['message']['content']
            
            # 保存AI响应到会话历史
            g.chat_history[session_id].append({'role': 'assistant', 'content': ai_response})
            
            return jsonify({'response': ai_response})
        else:
            # API响应异常
            error_message = response_json.get('error', {}).get('message', '未知错误')
            print(f"API错误: {error_message}")
            return jsonify({'response': f"抱歉，我无法处理您的请求。错误: {error_message}"})
    
    except Exception as e:
        print(f"处理请求时出错: {str(e)}")
        return jsonify({'response': "抱歉，处理您的请求时出现了问题。请稍后再试。"}), 500 